__author__ = 'grigory'

import sympy as smp
import copy
import scipy as sc


def ch(alpha, index, is_plus):
    beta = copy.copy(alpha)
    if is_plus:
        if beta[2 * index + 1] > 0:
            beta[2 * index + 1] -= 1
            beta[2 * index] += 1
            return beta
    else:
        if beta[2 * index] > 0:
            beta[2 * index + 1] += 1
            beta[2 * index] -= 1
            return beta
    return []


def minus(alpha, beta):
    gamma = []
    for i in range(len(alpha)):
        gamma += [alpha[i] - beta[i]]
    return gamma


def get_all_less(alpha):
    if len(alpha) == 1:
        return [[i] for i in range(alpha[0] + 1)]
    else:
        res = []
        beta = get_all_less(alpha[1:])
        for element in range(alpha[0] + 1):
            res += [[element] + b for b in beta]
    return res


'''
iter = iteration number
'''


def next_iteration_for_alpha_r(pols, iter, variables, variable, k, d, A, alpha):
    # 1st sum
    sum1 = smp.sympify("0")
    if tuple([variable] + alpha) in pols[iter - 1]:
        pol = pols[iter - 1][tuple([variable] + alpha)]
        for i in range(k):
            Au = "("
            for j in range(d):
                Au += str(A[i][j]) + "*" + variables[j] + "+"
            Au = Au[:-1]
            Au += ")"
            sum1 += smp.diff(pol, variables[i]) * smp.sympify(Au)

    #2nd sum
    sum2 = smp.sympify("0")
    betas = get_all_less(alpha)
    for beta in betas:
        for i in range(k):
            if tuple([variable] + beta) in pols[iter - 1]:
                if tuple([variables[i]] + minus(alpha, beta)) in pols[0]:
                    pol1 = pols[iter - 1][tuple([variable] + beta)]
                    pol2 = pols[0][tuple([variables[i]] + minus(alpha, beta))]
                    sum2 += smp.diff(pol1, variables[i]) * smp.sympify(pol2)

    #3rd sum
    sum3 = smp.sympify("0")
    for i in range(0, d - k):
        Au = "("
        for j in range(d):
            Au += str(A[k + i][j]) + "*" + variables[j] + "+"
        Au = Au[:-1]
        Au += ")"
        if tuple([variable] + ch(alpha, i, True)) in pols[iter - 1]:
            beta = ch(alpha, i, True)
            pol = pols[iter - 1][tuple([variable] + beta)]
            sum3 += smp.sympify(str(beta[2 * i])) * smp.sympify(pol) * smp.sympify(Au)
        if tuple([variable] + ch(alpha, i, False)) in pols[iter - 1]:
            beta = ch(alpha, i, False)
            pol = pols[iter - 1][tuple([variable] + beta)]
            sum3 -= smp.sympify(str(beta[2 * i + 1])) * smp.sympify(pol) * smp.sympify(Au)

    #4th sum
    sum4 = smp.sympify("0")
    for beta in betas:
        for i in range(d - k):
            if tuple([variable] + ch(beta, i, True)) in pols[iter - 1]:
                pol1 = pols[iter - 1][tuple([variable] + ch(beta, i, True))]
                if tuple([variables[k + i]] + minus(alpha, beta)) in pols[0]:
                    pol2 = pols[0][tuple([variables[k + i]] + minus(alpha, beta))]
                    sum4 += smp.sympify(pol1) * smp.sympify(pol2) * smp.sympify(ch(beta, i, True)[2 * i])
            if tuple([variable] + ch(beta, i, False)) in pols[iter - 1]:
                pol1 = pols[iter - 1][tuple([variable] + ch(beta, i, False))]
                if tuple([variables[k + i]] + minus(alpha, beta)) in pols[0]:
                    pol2 = pols[0][tuple([variables[k + i]] + minus(alpha, beta))]
                    sum4 -= smp.sympify(pol1) * smp.sympify(pol2) * smp.sympify(ch(beta, i, False)[2 * i + 1])

    return sum1 + sum2 + sum3 + sum4


def next_iteration_for_system(pols, iteration, variables, k, d, A, max_initial_degs):
    topdeg = []
    for i in range(len(max_initial_degs)):
        topdeg += [max_initial_degs[i] * (iteration + 1)]
    possible_degs = get_all_less(topdeg)
    polsForIteration = dict([])

    for possible_deg in possible_degs:
        for variable in variables:
            polsForIteration[tuple([variable] + possible_deg)] = next_iteration_for_alpha_r(pols, iteration, variables,
                                                                                            variable, k, d, A,
                                                                                            possible_deg)
    pols += [polsForIteration]
    return pols


def matlab_f_n(pols, n, variables):
    pol = pols[n]
    res = dict([])
    for (t, p) in pol.iteritems():
        if str(p) != "0":
            if t[0] not in res:
                res[t[0]] = ""
            res[t[0]] += "(" + str(p) + ")*" + form_matlab_line(t, variables) + "+"
    for v in variables:
        res[v] = res[v][:-1]
        res[v] = res[v].replace("**", "^")
    return res


def matlab_u_n(variables, u_n, f_n, A):
    res = dict([])
    for i in range(len(variables)):
        v = variables[i]
        res[v] = "-("
        for j in range(len(variables)):
            res[v] += "(" + str(A[i][j]) + ")*" + u_n[variables[j]] + "+"
        res[v] = res[v][:-1]
        res[v] += ")-"
        res[v] += f_n[v]
    return res


# assume A to be diagonal
def matlab_s_n(n, delta, k, q, A):
    s_n = []
    for qq in range(q):
        s = 0
        s += ((-1) ** (n + 1)) * ((-float(A[qq][qq]) + k) ** (-n - 1))
        m = sc.exp(-delta * (float(A[qq][qq]) + k))
        for nn in range(n + 1):
            s += ((-1) ** nn) * (delta ** (n - nn)) * ((-float(A[qq][qq]) + k) ** (-nn - 1)) * m
        s_n += [s]
    return s_n


def form_matlab_line(t, variables):
    d = len(variables)
    k = d - (len(t) - 1) / 2

    res = ""
    for i in range(d - k):
        res += "(sin(" + variables[k + i] + "))^(" + str(t[2 * (i + 1) - 1]) + ")*(cos(" + variables[
            k + i] + "))^(" + str(t[2 * (i + 1)]) + ")*"
    res = res[:-1]
    return res


def write_to_m_file(pols, iteration, variables, q, k, delta, A):
    f_ns = []
    u_0 = dict([])
    for v in variables:
        u_0[v] = v
    u_ns = [u_0]

    s_ns = []

    for i in range(iteration + 1):
        f_ns += [matlab_f_n(pols, i, variables)]
    for i in range(1, iteration + 1):
        u_ns += [matlab_u_n(variables, u_ns[i - 1], f_ns[i - 1], A)]
    for i in range(iteration + 1):
        s_ns += [matlab_s_n(i, delta, k, q, A)]

    res = dict([])
    for i in range(q):
        res[variables[i]] = variables[i] + "-("
        for j in range(iteration + 1):
            res[variables[i]] += "(" + str(s_ns[j][i]) + ")*(" + f_ns[j][variables[i]] + "-" + str(k) + "*(" + u_ns[j][
                variables[i]] + ")) + "
        res[variables[i]] = res[variables[i]][:-2]
        res[variables[i]] = res[variables[i]].replace("**", "^")

    final_str = ""
    for i in range(q):
        final_str += "(" + res[variables[i]] + ")) + "
    final_str = final_str[:-2]
    fl = open("polynom" + str(iteration) + ".m", "w")
    line = "function answer = polynom" + str(iteration) + "("
    for v in variables:
        line += v
        line += ", "
    line = line[:-2]
    line += ")\n"
    fl.write(line)
    line = "answer = "
    line += str(smp.simplify(smp.sympify(final_str))).replace("**", "^")
    line += ";\n"
    fl.write(line)
    fl.write("end\n")
    fl.close()

    
def main():
    #open config file
    f = open("config", "r")

    #read variables name
    line = f.next()
    variables = line.strip().split(" ")
    d = len(variables)

    #read k
    line = f.next()
    k = int(line.strip())

    #read matrix A
    A = []
    for i in range(d):
        a = f.next().strip().split(" ")
        A += [a]

    #maxdegrees for possible values estimation
    max_initial_degs = [1] * (2 * (d - k))

    #deal with lines
    pols = []
    polsForIteration = dict([])
    for line in f:
        line = line.strip()
        chunks = line.split(" ")
        for i in range(1, len(chunks) - 1):
            chunks[i] = int(chunks[i])
            if max_initial_degs[i - 1] < chunks[i]:
                max_initial_degs[i - 1] = chunks[i]
        polsForIteration[tuple(chunks[:-1])] = chunks[-1]
    pols += [polsForIteration]
    f.close()

    for xxx in pols:
        print xxx

    for iteration in range(7):
        pols = next_iteration_for_system(pols, iteration + 1, variables, k, d, A, max_initial_degs)
        print(str(iteration) + " completed")
    #print matlab_f_n(pols, 2, variables)

    for iteration in range(8):
        write_to_m_file(pols, iteration, variables, 2, 10, 0.001, A)
        print str(iteration) + " file done"

main()
